import matplotlib

from captcha.image import ImageCaptcha
import matplotlib.pyplot as plt
import numpy as np
import random
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"]='2' # 只显示 warning 和 Error

from tensorflow.python.keras.utils import Sequence
import string
characters = string.digits + string.ascii_uppercase


width, height, n_len, n_class = 90, 30, 4, len(characters)
# width, height, n_len, n_class = 170, 80, 4, len(characters)



from tensorflow.python.keras.models import *
from tensorflow.python.keras.layers import *

input_tensor = Input((height, width, 3))
x = input_tensor
# for i, n_cnn in enumerate([2, 2, 2, 2]):
for i, n_cnn in enumerate([2, 2, 2, 2]):
    for j in range(n_cnn):
        x = Conv2D(32*2**min(i, 3), kernel_size=3, padding='same', kernel_initializer='he_uniform')(x)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
    x = MaxPooling2D(2)(x)

x = Flatten()(x)
x = [Dense(n_class, activation='softmax', name='c%d'%(i+1))(x) for i in range(n_len)]
model = Model(inputs=input_tensor, outputs=x)
model.summary()
# exit()

from tensorflow.python.keras.utils import plot_model
from IPython.display import Image

# plot_model(model, to_file='cnn.png', show_shapes=True)
# Image('cnn.png')
#


from tensorflow.python.keras.callbacks import EarlyStopping, CSVLogger, ModelCheckpoint
from tensorflow.python.keras.optimizers import *
from tf_seq_generator import BaseSequence





if __name__=='__main__':
    print(characters)
    # img_paths_train = [r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train1'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train2'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train3'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train4'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train5'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train6'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train7'
    #                     ,r'/data/code_pic/szjj_gzh.261.szjjgateway_01.train8'
    #                    ]
    # img_paths_test = [r'/data/code_pic/szjj_gzh.261.szjjgateway_01.test']
    img_paths_train = [r'E:\szjj_gzh.261.szjjgateway_01.train']
    img_paths_test = [r'E:\szjj_gzh.261.szjjgateway_01.test']
    train_data = BaseSequence(paths=img_paths_train, characters=characters, batch_size=128, steps=300, n_len=n_len, width=width, height=height)
    valid_data = BaseSequence(paths=img_paths_test, characters=characters, batch_size=128, steps=30, n_len=n_len, width=width, height=height)
    # train_data = BaseSequence(characters, batch_size=128, steps=1000)
    # valid_data = BaseSequence(characters, batch_size=128, steps=20)
    callbacks = [EarlyStopping(patience=3), CSVLogger('cnn_szjj.csv'),
                 ModelCheckpoint('cnn_best_szjj.h5', save_best_only=True)]

    model.compile(loss='categorical_crossentropy',
                  optimizer=Adam(1e-3, amsgrad=True),
                  metrics=['accuracy'])
    model.fit_generator(train_data, epochs=100, validation_data=valid_data, workers=1, use_multiprocessing=True,
                        callbacks=callbacks)

    ### 载入最好的模型继续训练一会
    # 为了让模型充分训练，我们可以载入之前最好的模型权值，然后降低学习率为原来的十分之一，继续训练，这样可以让模型收敛得更好。

    # model.load_weights('cnn_best_win.h5')
    # callbacks = [EarlyStopping(patience=3), CSVLogger('cnn_win.csv', append=True),
    #              ModelCheckpoint('cnn_best_win.h5', save_best_only=True)]
    #
    # model.compile(loss='categorical_crossentropy',
    #               optimizer=Adam(1e-4, amsgrad=True),
    #               metrics=['accuracy'])
    # model.fit_generator(train_data, epochs=100, validation_data=valid_data, workers=4, use_multiprocessing=True,
    #                     callbacks=callbacks)

